#![allow(dead_code)]

use clap::{command, Parser, ValueEnum};
use framebuffer::Framebuffer;
use image::io::Reader as ImageReader;
use image::{GenericImageView, Rgb, RgbImage};

//Algorithm copied from:
//https://en.wikipedia.org/wiki/Mandelbrot_set

fn gen_mandelbrot_set(fb: &Framebuffer) -> Vec<u8> {
    let w = fb.var_screen_info.xres;
    let h = fb.var_screen_info.yres;
    let line_length = fb.fix_screen_info.line_length;
    let bytespp = fb.var_screen_info.bits_per_pixel / 8;
    let mut frame = vec![0u8; (line_length * h) as usize];
    for (r, line) in frame.chunks_mut(line_length as usize).enumerate() {
        for (c, p) in line.chunks_mut(bytespp as usize).enumerate() {
            let x0 = (c as f32 / w as f32) * 3.5 - 2.5;
            let y0 = (r as f32 / h as f32) * 2.0 - 1.0;

            let mut it = 0;
            let max_it = 200;

            let mut x = 0.0;
            let mut y = 0.0;

            while x * x + y * y < 4.0 && it < max_it {
                let xtemp = x * x - y * y + x0;
                y = 2.0 * x * y + y0;
                x = xtemp;
                it += 1;
            }

            p[0] = (125.0 * (it as f32 / max_it as f32)) as u8;
            p[1] = (255.0 * (it as f32 / max_it as f32)) as u8;
            p[2] = (75.0 * (it as f32 / max_it as f32)) as u8;
        }
    }
    frame
}

fn gen_clear_screen(fb: &Framebuffer) -> Vec<u8> {
    // let w = fb.var_screen_info.xres;
    let h = fb.var_screen_info.yres;
    let line_length = fb.fix_screen_info.line_length;
    let bytespp = fb.var_screen_info.bits_per_pixel / 8;
    let mut frame = vec![0u8; (line_length * h) as usize];
    for line in frame.chunks_mut(line_length as usize) {
        for p in line.chunks_mut(bytespp as usize) {
            p[0] = 0;
            p[1] = 0;
            p[2] = 0;
        }
    }
    frame
}

// png file is ok, while jpg failed.
fn gen_frame_through_imagefile(fb: &Framebuffer, input_filename: &String) -> Vec<u8> {
    let img = ImageReader::open(input_filename).unwrap().decode().unwrap();
    println!("img w:{},h:{}", img.width(), img.height());
    // let w = fb.var_screen_info.xres;
    let h = fb.var_screen_info.yres;
    let line_length = fb.fix_screen_info.line_length;
    let bytespp = fb.var_screen_info.bits_per_pixel / 8;
    let mut frame = vec![0u8; (line_length * h) as usize];
    for (y, line) in frame.chunks_mut(line_length as usize).enumerate() {
        for (x, p) in line.chunks_mut(bytespp as usize).enumerate() {
            let x = x as u32;
            let y = y as u32;
            if x < img.width() && y < img.height() {
                let rgba = img.get_pixel(x, y);
                p[0] = rgba[0];
                p[1] = rgba[1];
                p[2] = rgba[2];
            } else {
                p[0] = 0;
                p[1] = 0;
                p[2] = 0;
            }
        }
    }
    frame
}
fn save_framebuffer(fb: &Framebuffer, out_filename: &String) {
    let w = fb.var_screen_info.xres;
    let h = fb.var_screen_info.yres;
    let line_length = fb.fix_screen_info.line_length;
    let bytespp = fb.var_screen_info.bits_per_pixel / 8;
    let frame = fb.read_frame();
    let mut img = RgbImage::new(w, h);
    for (y, line) in frame.chunks(line_length as usize).enumerate() {
        for (x, p) in line.chunks(bytespp as usize).enumerate() {
            img.put_pixel(x as u32, y as u32, Rgb([p[2], p[1], p[0]]));
        }
    }
    img.save_with_format(out_filename, image::ImageFormat::Png)
        .unwrap();
}

#[derive(Parser)]
#[command(version, author, about, long_about = None)]
struct Cli {
    #[arg(short, long)]
    mode: Option<Mode>,
    #[arg(short, long)]
    filename: Option<String>,
    #[arg(short, long)]
    fb_dev_name: Option<String>,
}

#[derive(Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
enum Mode {
    /// read from frame buffer and save to image file
    #[default]
    Read,
    /// gen test image and write to frame buffer
    Write,
    /// clear screen
    Clear,
    /// get frame buffer detail info
    Info,
}

fn main() {
    let cli = Cli::parse();

    let fb_dev_name = cli.fb_dev_name.unwrap_or(String::from("/dev/fb0"));
    let mut filename = cli.filename.unwrap_or(String::from(""));
    let mode = cli.mode.unwrap_or_default();

    match mode {
        Mode::Read => {
            let framebuffer = Framebuffer::new(fb_dev_name).unwrap();
            if filename.is_empty() {
                filename = String::from("output.png");
            }
            println!("read screen and save to {}", filename);
            save_framebuffer(&framebuffer, &filename);
        }
        Mode::Write => {
            let mut framebuffer = Framebuffer::new(fb_dev_name).unwrap();
            if filename.is_empty() {
                println!("write screen with mandelbrot...");
                let frame = gen_mandelbrot_set(&framebuffer);
                framebuffer.write_frame(&frame);
            } else {
                println!("write screen with file {}", filename);
                let frame = gen_frame_through_imagefile(&framebuffer, &filename);
                framebuffer.write_frame(&frame);
            }
        }
        Mode::Clear => {
            println!("clear screen...");
            let mut framebuffer = Framebuffer::new(fb_dev_name).unwrap();
            let frame = gen_clear_screen(&framebuffer);
            framebuffer.write_frame(&frame);
        }
        Mode::Info => {
            println!("get screen...");
            let framebuffer = Framebuffer::new(fb_dev_name).unwrap();
            println!("var_screen_info:{:#?}", framebuffer.var_screen_info);
            println!("fix_screen_info:{:#?}", framebuffer.fix_screen_info);
            println!("frame len:{}", framebuffer.frame.len());
        }
    }

    println!("press to exit...");
    std::io::stdin().read_line(&mut String::new()).unwrap();
}
